import numpy as np
import cv2
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Flatten, Dense

# 加载原始图像
image = cv2.imread('image.jpg')

# 加载水印图像
watermark = cv2.imread('watermark.png')

# 预处理水印图像（可以根据需要进行调整）
watermark = cv2.resize(watermark, (image.shape[1], image.shape[0]))

# 创建水印嵌入模型
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=image.shape))
model.add(Flatten())
model.add(Dense(1024, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

# 编译模型
model.compile(loss='binary_crossentropy', optimizer='adam')

# 训练模型以嵌入水印
model.fit(image, watermark, epochs=10)

# 嵌入水印到图像
watermarked_image = model.predict(image)

# 显示原始图像、水印和水印后的图像
cv2.imshow('Original Image', image)
cv2.imshow('Watermark', watermark)
cv2.imshow('Watermarked Image', watermarked_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 提取水印
extracted_watermark = model.predict(watermarked_image)

# 显示提取的水印
cv2.imshow('Extracted Watermark', extracted_watermark)
cv2.waitKey(0)
cv2.destroyAllWindows()
